# -*- coding: utf-8 -*-
# @Time    : 2025/7/11 14:49
# @Author  : XinHuang
# @Email   : SA24231052@mail.ustc.edu.cn
# @File    : model
# @Software: PyCharm
import numpy as np
import matplotlib.pyplot as plt
import logging


class FeedBackSystem:
    def __init__(self, R, omega0, Q, omega, T_0):
        self.R = R
        self.omega0 = omega0
        self.Q = Q
        self.omega = omega
        self.T_0 = T_0

        self.transverse_impedance = None
        self.transverse_impedance_fb = None
        self.longitude_impedance = None
        self.transverse_impedance_OTFB = None
        self.h_fb = None
        self.h_OTFB = None

    def calculate_impedance(self) -> None:
        """
        计算开环阻抗，公式为：
                      w0*R                     w
        Z_x(w) = ------------------ , Z_// = ---- Z_x(w)
                 w[1+iQ(w0/w-w/w0)]          w0*R

        :param self.omega: numpy数组
        """
        self.transverse_impedance = self.omega0 * self.R / (
                self.omega * (1 + 1j * self.Q * (self.omega0 / self.omega - self.omega / self.omega0)))
        self.longitude_impedance = self.omega / (self.omega0 * self.R) * self.transverse_impedance
        logging.info(f"Open loop cavity max abs impedance: {np.max(np.abs(self.transverse_impedance))}")

    def add_feedback(self, G_fb, tao_d):
        """"""
        G_opt = self.Q / (self.omega0 * tao_d)
        G_fb = G_opt
        # delta_omega = self.omega - self.omega0 * self.omega / np.abs(self.omega)
        logging.info(f'G_fb={G_fb:.4e}')
        delta_omega = self.omega - self.omega0
        self.h_fb = G_fb * np.exp(-1j * delta_omega * tao_d) * self.longitude_impedance
        self.transverse_impedance_fb = self.transverse_impedance / (1 - self.h_fb)  # 注意，论文里面写错了，应该是减号
        logging.info(f"Closed loop cavity max abs impedance: {np.max(np.abs(self.transverse_impedance_fb))}")

    def add_OTFB(self, G_c, alpha_c, mu_x, tao_c, m):
        """"""
        logging.info(f'G_c={G_c:.4e}, T_0={self.T_0}')
        delta_omega = self.omega - self.omega0
        if m == 0:
            part1 = G_c * (1 - alpha_c) * np.exp(1j * 2 * np.pi * mu_x) * np.exp(-1j * (self.T_0 - tao_c) * delta_omega)
            part2 = 1 - alpha_c * np.exp(1j * 2 * np.pi * mu_x) * np.exp(-1j * self.T_0 * delta_omega)
            self.h_OTFB = part1 / part2
        else:
            raise "This part of the work hasn't been done yet..."
        # self.transverse_impedance_OTFB = self.transverse_impedance / (1 + self.h_OTFB)
        # self.transverse_impedance_OTFB = self.transverse_impedance / (1 + self.h_OTFB) / (1 - self.h_fb)
        self.transverse_impedance_OTFB = self.transverse_impedance / (1 - self.h_fb * (1 + self.h_OTFB))

    def plot_results(self, ):
        """"""
        fontdict = {
            'weight': 'bold',
            'size': 14,
            'fontproperties': 'Times New Roman'
        }
        figdict = {
            # 刻度方向设置为 "in"，表示刻度向内
            'xtick.direction': 'in',
            'ytick.direction': 'in',
            # 刻度标签字体大小设置为 12，符合学术论文中常见的字号要求
            'xtick.labelsize': 12,
            'ytick.labelsize': 12,
            # 字体族设置为 serif，学术论文中通常使用 serif 字体（如 Times New Roman）以增强可读性
            'font.family': 'serif',
            # 指定字体为 Times New Roman，这是学术论文中常用的字体之一
            'font.serif': ['Times New Roman'],
            # 设置标题字体大小为 14，稍大于刻度标签字体，以突出显示
            'axes.titlesize': 14,
            # 设置轴标签字体大小为 12，与刻度标签字体大小一致
            'axes.labelsize': 12,
            # 设置图例字体大小为 10，确保图例文字清晰可读但不喧宾夺主
            'legend.fontsize': 10,
            # 设置图形尺寸为宽 6 英寸、高 4 英寸，这是一个常见的学术图表尺寸，既不太大也不太小
            'figure.figsize': (6, 4),
            # 设置图形分辨率为 600 DPI，这是印刷出版物中常用的标准分辨率，能够保证图片质量
            'figure.dpi': 300,
            # 设置底部边距为 0.15，确保标题和轴标签有足够的空间显示
            'figure.subplot.bottom': 0.15,
            # 设置左侧边距为 0.15，确保 y 轴标签和刻度有足够的空间显示
            'figure.subplot.left': 0.15,
            # 设置右侧边距为 0.95，避免右侧内容超出图形边界
            'figure.subplot.right': 0.95,
            # 设置顶部边距为 0.95，避免标题超出图形边界
            'figure.subplot.top': 0.95,
            # 设置线条宽度为 1.0，确保线条清晰可见但不过于粗壮
            'lines.linewidth': 1.0,
            # 设置标记大小为 4，使数据点标记清晰可辨
            'lines.markersize': 4,
            # 设置网格线样式为虚线，增强图表的可读性同时不干扰数据线条
            'grid.linestyle': '--',
            # 设置网格线宽度为 0.5，使网格线不喧宾夺主
            'grid.linewidth': 0.5
        }
        plt.rcParams.update(figdict)
        x_data = (self.omega - self.omega0) / 2 / np.pi / 1000
        plt.figure(1)
        plt.title('abs transverse impedance', fontdict=fontdict)
        plt.semilogy(x_data, np.abs(self.transverse_impedance), label='Open loop Cavity')
        plt.semilogy(x_data, np.abs(self.transverse_impedance_fb), label='Closed Loop')
        plt.semilogy(x_data, np.abs(self.transverse_impedance_OTFB), label='Closed Loop + OTFB')
        # plt.semilogy(x_data, np.abs(np.real(self.transverse_impedance)), label='Open loop Cavity')
        # plt.semilogy(x_data, np.abs(np.real(self.transverse_impedance_fb)), label='Closed Loop')
        # plt.semilogy(x_data, np.abs(np.real(self.transverse_impedance_OTFB)), label='Closed Loop + OTFB')
        plt.xlabel(r"$f-f_0$/(kHz)", fontdict=fontdict)
        plt.ylabel("impedance", fontdict=fontdict)
        plt.xlim([-8e2, 8e2])
        # plt.ylim([1e5, 1e11])
        plt.legend()
        plt.tight_layout()
        plt.savefig('fig/Fig1.png')
        plt.close()

        plt.figure(2)
        plt.subplot(3, 1, 1)
        plt.title('abs h_fb', fontdict=fontdict)
        plt.semilogy(x_data, np.abs(self.h_fb))
        plt.xlim([-8e2, 8e2])
        plt.subplot(3, 1, 2)
        plt.title('real h_fb', fontdict=fontdict)
        plt.semilogy(x_data, np.real(self.h_fb))
        plt.xlim([-8e2, 8e2])
        plt.subplot(3, 1, 3)
        plt.title('imag h_fb', fontdict=fontdict)
        plt.plot(x_data, np.imag(self.h_fb))
        plt.xlabel(r"$f-f_0$/(kHz)", fontdict=fontdict)
        plt.xlim([-8e2, 8e2])
        plt.tight_layout()
        plt.savefig('fig/Fig2.png')
        plt.close()

        # plt.figure(3)
        # plt.title('real transverse impedance', fontdict=fontdict)
        # plt.plot(x_data, 2 * np.real(self.transverse_impedance), label='Open loop Cavity')
        # plt.plot(x_data, 2 * np.real(self.transverse_impedance_fb), label='Closed Loop')
        # plt.xlabel(r"$f-f_0$/(kHz)", fontdict=fontdict)
        # plt.ylabel("impedance", fontdict=fontdict)
        # plt.xlim([-8e6, 8e6])
        # plt.tight_layout()

        # plt.figure(4)
        # plt.title('imag transverse impedance', fontdict=fontdict)
        # plt.plot(x_data, 2 * np.imag(self.transverse_impedance), label='Open loop Cavity')
        # plt.plot(x_data, 2 * np.imag(self.transverse_impedance_fb), label='Closed Loop')
        # plt.xlabel(r"$f-f_0$/(kHz)", fontdict=fontdict)
        # plt.ylabel("impedance", fontdict=fontdict)
        # plt.xlim([-8e6, 8e6])
        # plt.tight_layout()

        plt.figure(5)
        plt.title(r'abs $Z_{//}$', fontdict=fontdict)
        plt.semilogy(x_data, np.abs(self.longitude_impedance))
        plt.xlabel(r"$f-f_0$/(kHz)", fontdict=fontdict)
        plt.xlim([-8e2, 8e2])
        plt.tight_layout()
        plt.savefig('fig/Fig5.png')
        plt.close()

        plt.figure(6)
        plt.title(r'$H_{OTFB}$', fontdict=fontdict)
        plt.plot(x_data, np.abs(1+self.h_OTFB), label ='abs')
        plt.plot(x_data, np.real(1+self.h_OTFB), label='real')
        plt.plot(x_data, np.imag(1+self.h_OTFB), label='imag')
        plt.xlabel(r"$f-f_0$/(kHz)", fontdict=fontdict)
        plt.legend()
        plt.xlim([-8e2, 8e2])
        plt.tight_layout()
        plt.savefig('fig/Fig6.png')
        plt.close()

        # plt.show(block=True)
